from fastapi import FastAPI, UploadFile, File, Form
from fastapi.responses import StreamingResponse
from rembg import remove, new_session
import cv2
import numpy as np
from PIL import Image
import io
import os
import glob
import logging
from pathlib import Path
import torch

device = "cuda" if torch.cuda.is_available() else "cpu"

# 配置日志，详细记录颜色和服装处理过程
logging.basicConfig(level=logging.DEBUG, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

app = FastAPI(title="增强版证件照制作API")

# 预加载人像分割模型
session = new_session("u2net_human_seg", device=device)

# 服装资源配置（使用绝对路径避免相对路径问题）
CLOTHES_DIR = Path(__file__).parent / "id_photo_clothes"
try:
    CLOTHES_DIR.mkdir(parents=True, exist_ok=True)
    logger.info(f"服装文件夹已创建或存在：{CLOTHES_DIR.resolve()}")
except Exception as e:
    logger.error(f"创建服装文件夹失败：{str(e)}")

# 重新实现服装加载逻辑（确保正确识别PNG透明图）
def load_clothes_options():
    """加载并验证服装选项"""
    try:
        # 只加载PNG文件
        clothes_files = list(CLOTHES_DIR.glob("*.png"))
        logger.debug(f"找到{len(clothes_files)}个PNG文件在服装文件夹")
        
        valid_clothes = []
        for file in clothes_files:
            try:
                # 验证是否为透明图（RGBA模式）
                with Image.open(file) as img:
                    if img.mode == "RGBA":
                        valid_clothes.append(file)
                        logger.debug(f"有效服装文件：{file.name}")
                    else:
                        logger.warning(f"跳过非透明图：{file.name}（模式：{img.mode}）")
            except Exception as e:
                logger.error(f"验证服装文件{file.name}失败：{str(e)}")
        
        return [f.name for f in valid_clothes]
    except Exception as e:
        logger.error(f"加载服装选项失败：{str(e)}")
        return []


@app.post("/make_id_photo")
async def make_id_photo(
    file: UploadFile = File(...),
    bg_color: str = Form("white"),  # 支持 white/red/blue/#RRGGBB
    width: int = Form(295),
    height: int = Form(413),
    beauty_level: float = Form(0.3),
    clothes: str = Form(None)
):
    logger.info(f"接收到颜色参数：{bg_color}")
    
    # 1. 读取上传图片
    try:
        contents = await file.read()
        img = Image.open(io.BytesIO(contents)).convert("RGBA")
        logger.debug(f"成功读取图片：{file.filename}，尺寸：{img.size}")
    except Exception as e:
        logger.error(f"读取图片失败：{str(e)}")
        return {"error": f"读取图片失败：{str(e)}"}

    # 2. 人像抠图
    try:
        img_removed = remove(
            np.array(img),
            session=session,
            alpha_matting=True,
            alpha_matting_foreground_threshold=240,
            alpha_matting_background_threshold=10
        )
        img_cropped = Image.fromarray(img_removed).convert("RGBA")
        logger.debug("人像抠图完成")
    except Exception as e:
        logger.error(f"人像抠图失败：{str(e)}")
        return {"error": f"人像抠图失败：{str(e)}"}

    # 3. 应用美颜
    try:
        beauty_level = max(0, min(1, beauty_level))
        # 简单的美颜实现（避免影响背景色）
        img_np = np.array(img_cropped)
        if beauty_level > 0:
            # 只对人脸区域（非透明部分）应用美颜
            mask = img_np[:, :, 3] / 255.0
            blurred = cv2.GaussianBlur(img_np, (15, 15), 0)
            # 结合原图和模糊图，根据美颜强度调整
            img_np = (img_np * (1 - beauty_level) + blurred * beauty_level * mask[:, :, np.newaxis]).astype(np.uint8)
        img_beautified = Image.fromarray(img_np).convert("RGBA")
        logger.debug(f"美颜处理完成，强度：{beauty_level}")
    except Exception as e:
        logger.error(f"美颜处理失败：{str(e)}")
        img_beautified = img_cropped  # 失败时使用原图

    # 4. 叠加服装
    try:
        if clothes and clothes != "不使用服装":
            clothes_path = CLOTHES_DIR / clothes
            if clothes_path.exists():
                with Image.open(clothes_path) as clothes_img:
                    if clothes_img.mode != "RGBA":
                        logger.warning(f"服装{clothes}不是透明图，跳过叠加")
                    else:
                        # 调整服装大小以适应人像
                        img_width, img_height = img_beautified.size
                        person_height = int(img_height * 0.7)
                        clothes_width = int(clothes_img.size[0] * (person_height / clothes_img.size[1]))
                        clothes_resized = clothes_img.resize((clothes_width, person_height), Image.LANCZOS)
                        
                        # 计算放置位置（居中）
                        x = (img_width - clothes_width) // 2
                        y = int(img_height * 0.2)
                        
                        # 叠加服装
                        result = Image.new("RGBA", img_beautified.size)
                        result.paste(img_beautified, (0, 0))
                        result.paste(clothes_resized, (x, y), mask=clothes_resized.split()[3])
                        img_beautified = result
                        logger.debug(f"成功叠加服装：{clothes}")
            else:
                logger.warning(f"服装文件不存在：{clothes}")
    except Exception as e:
        logger.error(f"叠加服装失败：{str(e)}")

    # 5. 处理背景颜色（核心修复部分）
    try:
        # 预设标准证件照颜色（使用行业标准值）
        bg_map = {
            "white": (255, 255, 255, 255),  # 白色背景
            "red": (255, 0, 0, 255),        # 红色背景
            "blue": (67, 142, 219, 255)     # 标准证件照蓝色
        }

        # 解析颜色参数
        final_bg = bg_map["white"]  # 默认白色
        bg_color = bg_color.strip().lower()

        if bg_color in bg_map:
            final_bg = bg_map[bg_color]
            logger.info(f"使用预设颜色: {bg_color} -> RGBA{final_bg}")
        elif bg_color.startswith("#"):
            # 处理十六进制颜色
            try:
                hex_str = bg_color.lstrip("#")
                # 处理#RGB格式（如#F00）
                if len(hex_str) == 3:
                    hex_str = ''.join([c*2 for c in hex_str])
                # 处理#RRGGBB格式
                if len(hex_str) == 6:
                    r = int(hex_str[0:2], 16)
                    g = int(hex_str[2:4], 16)
                    b = int(hex_str[4:6], 16)
                    final_bg = (r, g, b, 255)
                    logger.info(f"解析自定义颜色: {bg_color} -> RGBA{final_bg}")
                else:
                    logger.warning(f"无效的十六进制颜色长度: {len(hex_str)}，使用默认白色")
            except Exception as e:
                logger.error(f"解析颜色失败: {bg_color}, 错误: {str(e)}，使用默认白色")
        else:
            logger.warning(f"未知颜色格式: {bg_color}, 使用默认白色")

        # 6. 叠加背景（核心修复：确保正确应用背景色）
        # 创建背景图层
        background = Image.new("RGBA", img_beautified.size, final_bg)
        # 合并图层（确保人像在背景之上）
        final_img = Image.alpha_composite(background, img_beautified)
        logger.debug("背景色应用完成")
    except Exception as e:
        logger.error(f"背景色处理失败：{str(e)}")
        final_img = img_beautified  # 失败时使用原图

    # 7. 调整尺寸
    try:
        final_img = final_img.convert("RGB")
        final_img.thumbnail((width, height), Image.LANCZOS)
        logger.debug(f"调整尺寸至: {width}x{height}")
    except Exception as e:
        logger.error(f"调整尺寸失败：{str(e)}")

    # 8. 输出
    try:
        img_byte_arr = io.BytesIO()
        final_img.save(img_byte_arr, format="JPEG", quality=95)
        img_byte_arr.seek(0)
        return StreamingResponse(img_byte_arr, media_type="image/jpeg")
    except Exception as e:
        logger.error(f"生成图片响应失败：{str(e)}")
        return {"error": f"生成图片失败：{str(e)}"}


@app.get("/clothes_options")
async def get_clothes_options():
    """获取可用服装选项（修复版本）"""
    try:
        clothes_list = load_clothes_options()
        clothes_path = str(CLOTHES_DIR.resolve())
        logger.info(f"返回服装选项：{len(clothes_list)}个，路径：{clothes_path}")
        return {
            "clothes_folder": clothes_path,
            "available_clothes": clothes_list,
            "total_count": len(clothes_list),
            "message": "将PNG透明图放入上述文件夹即可显示" if not clothes_list else ""
        }
    except Exception as e:
        error_msg = f"获取服装选项失败：{str(e)}"
        logger.error(error_msg)
        return {"error": error_msg}


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)